import csv
import ast
import numpy as np
import matplotlib.pyplot as plt
import torch
import math


# 读取 CSV 文件
data = {
    "ankle_angle": [],
    "ankle_torques": [],
    "left_imu_lin_acc": [],
    "left_imu_ang_vel": [],
    "right_imu_lin_acc": [],
    "right_imu_ang_vel": [],
    "left_ankle_height": [],
    "right_ankle_height": [],
    "joint_vel": [],
    "root_lin_vel_w": [],
}

with open("data_continus_walking_1dot25_extras_log.csv", "r") as f:
    reader = csv.reader(f)
    for row in reader:
        if len(row) < 9:
            continue
        # 使用 ast.literal_eval 安全地将字符串转为 Python 对象
        data["ankle_angle"].append(np.array(ast.literal_eval(row[0])).flatten())
        data["ankle_torques"].append(np.array(ast.literal_eval(row[1])).flatten())
        data["left_imu_lin_acc"].append(np.array(ast.literal_eval(row[2])).flatten())
        data["left_imu_ang_vel"].append(np.array(ast.literal_eval(row[3])).flatten())
        data["right_imu_lin_acc"].append(np.array(ast.literal_eval(row[4])).flatten())
        data["right_imu_ang_vel"].append(np.array(ast.literal_eval(row[5])).flatten())
        data["left_ankle_height"].append(ast.literal_eval(row[6]))
        data["right_ankle_height"].append(ast.literal_eval(row[7]))
        data["joint_vel"].append(np.array(ast.literal_eval(row[8])).flatten())
        data["root_lin_vel_w"].append(np.array(ast.literal_eval(row[9])).flatten())

# 转为 numpy 数组
for k in data:
    data[k] = np.vstack(data[k]) if isinstance(data[k][0], np.ndarray) else np.array(data[k])


# 删除前 1000 个数据点
for k in data:
    data[k] = data[k][1000:]

joint_vel = data["joint_vel"]

# 拆分左、右机器人
# left_indices_all = []
# right_indices_all = []

# for i in range(0, 64, 2):  # 偶数列 -> 左
#     idx = np.where(joint_vel[:, i] > 6)[0]
#     idx 
#     if len(idx) > 0:
#         left_indices_all.append((i, idx))  # (机器人编号, 对应时间索引)
#     else :
#         print(f"机器人编号 {i} 没有出现过速度>6")

# for i in range(1, 64, 2):  # 奇数列 -> 右
#     idx = np.where(joint_vel[:, i] > 6)[0]
#     if len(idx) > 0:
#         right_indices_all.append((i, idx))
# print(f"左机器人共有 {len(left_indices_all)} 个机器人出现过速度>6")
# print(f"右机器人共有 {len(right_indices_all)} 个机器人出现过速度>6")


height_colomn, height_row = data["left_ankle_height"].shape
swing_standing_label = np.zeros((height_colomn, height_row))
left_ankle_heights = data["left_ankle_height"]
right_ankle_heights = data["right_ankle_height"]

print(height_row,height_colomn)



# 存储每一列的摆动相位起止索引
swing_phases = []

for i in range(height_row):
    stance_started = False
    stance_counter = 0
    stance_start_index = None
    swing_started = False
    swing_start_index = None
    swing_end_index = None

    for j in range(height_colomn):
        # 检测左脚进入 stance 相位（高度低于 0.02 且连续超过 15 个点）
        if left_ankle_heights[j, i] < 0.02:
            stance_counter += 1
        else:
            stance_counter = 0  # 不连续就重置计数

        if stance_counter > 15 and not stance_started:
            stance_start_index = j - 15  # 第一次连续低于 0.02 的起点
            stance_started = True

        # 检测左脚进入 swing 相位（高度高于右脚）
        if stance_started and not swing_started and left_ankle_heights[j, i] > right_ankle_heights[j, i]:
            if j <  20: 
                detect_error_index = 0
            else :
                detect_error_index = j - 20
            if np.any(swing_standing_label[detect_error_index:j, i] <0):
                stance_started = False
            else:
                swing_start_index = j
                swing_standing_label[j, i] = 0.15
                swing_started = True

        # 检测左脚摆动相位结束（高度再次低于右脚）
        if swing_started and left_ankle_heights[j, i] < right_ankle_heights[j, i] and not (left_ankle_heights[j, i]> 0.03 and right_ankle_heights[j, i]> 0.03) :
            swing_end_index = j
            swing_standing_label[j, i] = 0.2
            swing_started = False  # 摆动结束
            # 检测这个区间是否有异常
            if np.any(left_ankle_heights[swing_start_index:swing_end_index, i] > 0.125)or np.any(right_ankle_heights[swing_start_index:swing_end_index, i] > 0.125):
                swing_standing_label[swing_start_index:swing_end_index+1, i] = -0.1
                stance_started = False
        

    # # 保存本列的摆动起止索引
    # swing_phases.append({
    #     "stance_start": stance_start_index,
    #     "swing_start": swing_start_index,
    #     "swing_end": swing_end_index
    # })

# 将不记入的区间设置为异常
for i in range(height_row):
    record_index = 0
    error_data_label = False
    for j in range(height_colomn):
        if swing_standing_label[j, i] == 0.2:
            record_index = j
        if swing_standing_label[j, i] == -0.1:
            error_data_label = True
        if error_data_label and swing_standing_label[j, i] == 0.15:
            swing_standing_label[record_index+1:j-1, i] = -0.1
            error_data_label = False


# 初始化容器
left_ankle_angle = np.zeros((height_colomn, height_row))
right_ankle_angle = np.zeros((height_colomn, height_row))
left_ankle_torque = np.zeros((height_colomn, height_row))
right_ankle_torque = np.zeros((height_colomn, height_row))
left_ankle_velocity = np.zeros((height_colomn, height_row))
right_ankle_velocity = np.zeros((height_colomn, height_row))

# for robot_idx in range(height_colomn):
#     print(f"机器人编号 {robot_idx} 开始处理")

# 按每两个列分配给一个机器人
for robot_idx in range(height_row):


    left_col = 2 * robot_idx
    right_col = 2 * robot_idx + 1

    left_ankle_angle[:, robot_idx] =  data["ankle_angle"][:, left_col]
    right_ankle_angle[:, robot_idx] =  data["ankle_angle"][:, right_col]

    left_ankle_velocity[:, robot_idx] =  data["joint_vel"][:, left_col]
    right_ankle_velocity[:, robot_idx] =  data["joint_vel"][:, right_col]

    left_ankle_torque[:, robot_idx] =  data["ankle_torques"][:, left_col]
    right_ankle_torque[:, robot_idx] = data["ankle_torques"][:, right_col]



imu_colomn, imu_row = data["right_imu_lin_acc"].shape
print(imu_row,imu_colomn)


left_imu_ang_vel_x = np.zeros((height_colomn, height_row))
left_imu_ang_vel_y = np.zeros((height_colomn, height_row))
left_imu_ang_vel_z = np.zeros((height_colomn, height_row))

left_imu_lin_acc_x = np.zeros((height_colomn, height_row))
left_imu_lin_acc_y = np.zeros((height_colomn, height_row))
left_imu_lin_acc_z = np.zeros((height_colomn, height_row))



right_imu_ang_vel_x = np.zeros((height_colomn, height_row))
right_imu_ang_vel_y = np.zeros((height_colomn, height_row))
right_imu_ang_vel_z = np.zeros((height_colomn, height_row))

right_imu_lin_acc_x = np.zeros((height_colomn, height_row))
right_imu_lin_acc_y = np.zeros((height_colomn, height_row))
right_imu_lin_acc_z = np.zeros((height_colomn, height_row))

for robot_idx in range(height_row):

    x = 3 * robot_idx
    y = 3 * robot_idx + 1
    z = 3 * robot_idx + 2
    left_imu_ang_vel_x[:, robot_idx] = data["left_imu_ang_vel"][:, x]
    left_imu_ang_vel_y[:, robot_idx] = data["left_imu_ang_vel"][:, y]
    left_imu_ang_vel_z[:, robot_idx] = data["left_imu_ang_vel"][:, z]

    left_imu_lin_acc_x[:, robot_idx] = data["left_imu_lin_acc"][:, x]
    left_imu_lin_acc_y[:, robot_idx] = data["left_imu_lin_acc"][:, y]
    left_imu_lin_acc_z[:, robot_idx] = data["left_imu_lin_acc"][:, z]

    right_imu_ang_vel_x[:, robot_idx] = data["right_imu_ang_vel"][:, x]
    right_imu_ang_vel_y[:, robot_idx] = data["right_imu_ang_vel"][:, y]
    right_imu_ang_vel_z[:, robot_idx] = data["right_imu_ang_vel"][:, z]

    right_imu_lin_acc_x[:, robot_idx] = data["right_imu_lin_acc"][:, x]
    right_imu_lin_acc_y[:, robot_idx] = data["right_imu_lin_acc"][:, y]
    right_imu_lin_acc_z[:, robot_idx] = data["right_imu_lin_acc"][:, z]

    



dataset = []

for i in range(height_row):
    left_swing = False
    for j in range(height_colomn):
        if swing_standing_label[j, i] == 0.15 and not left_swing:
            # 左脚摆动开始
            left_swing_start_index = j
            left_swing = True

        elif swing_standing_label[j, i] == 0.2 and left_swing:
            # 左脚摆动结束
            left_swing_end_index = j
            left_swing = False

            # 计算中间点（转换点）
            center_index = math.ceil((float(left_swing_start_index) + float(left_swing_end_index)) / 2)

            if left_swing_start_index % 2 == 1 and left_swing_end_index % 2 == 0:
                left_swing_end_index += 1

            if left_swing_start_index % 2 == 0 and left_swing_end_index % 2 == 1:
                left_swing_end_index += 1
 
                
            # 生成 src：左脚的摆动段（T1）
            src = np.concatenate([
                left_ankle_angle[left_swing_start_index:center_index, i][:, None],
                left_ankle_velocity[left_swing_start_index:center_index, i][:, None],
                left_imu_ang_vel_x[left_swing_start_index:center_index, i][:, None],
                left_imu_ang_vel_y[left_swing_start_index:center_index, i][:, None],
                left_imu_ang_vel_z[left_swing_start_index:center_index, i][:, None],
                left_imu_lin_acc_x[left_swing_start_index:center_index, i][:, None],
                left_imu_lin_acc_y[left_swing_start_index:center_index, i][:, None],
                left_imu_lin_acc_z[left_swing_start_index:center_index, i][:, None],
                data["left_ankle_height"][left_swing_start_index:center_index, i][:, None],

            ], axis=1)

            tgt =left_ankle_torque[left_swing_start_index -3: center_index, i][:, None]
            dataset.append((src, tgt))


            # 添加一条数据 (src, tgt)
            dataset.append((src, tgt))

print(len(dataset))
src, tgt = dataset[2]
print("src shape:", src.shape)
print("tgt shape:", tgt.shape)
# 将 dataset 中的所有样本收集到一个列表
all_data = [(src, tgt) for src, tgt in dataset]

# 保存为 PyTorch 文件
torch.save(all_data, "dataset_1dot25_continus_walk.pt")



# 打印形状确认
for k, v in data.items():
    print(f"{k}: {v.shape}")

# ==============================
#      开始绘图
# ==============================
time = np.arange(len(data["ankle_angle"]))

plt.figure(figsize=(14, 10))

# # 1. 踝关节角度
# plt.subplot(3, 3, 1)
# plt.plot(time, data["ankle_angle"][:, 0], label="Left ankle angle")
# plt.plot(time, data["ankle_angle"][:, 1], label="Right ankle angle")
# plt.title("Ankle Angles")
# plt.legend()

# 2. 踝关节力矩
# plt.subplot(1, 1, 1)
# ax1 = plt.gca()  # 获取当前坐标轴

# # 第一个纵轴：踝关节力矩
# ax1.plot(time, data["ankle_torques"][:, 0], label="Left torque", color='b')
# ax1.plot(time, data["ankle_torques"][:, 1], label="Right torque", color='r')
# ax1.set_ylabel("Ankle Torque (Nm)")

# # 第二个纵轴：Swing/Standing
# ax2 = ax1.twinx()
# ax2.plot(time, swing_standing_label[:, 2], label="Swing/Standing", color='g', linestyle='--')
# ax2.set_ylabel("Swing/Standing")

# # 合并图例
# lines_1, labels_1 = ax1.get_legend_handles_labels()
# lines_2, labels_2 = ax2.get_legend_handles_labels()
# ax1.legend(lines_1 + lines_2, labels_1 + labels_2, loc="upper right")

# plt.title("Ankle Torques")

# # 3. 左右脚 IMU 加速度
# plt.subplot(3, 3, 3)
# plt.plot(time, data["left_imu_lin_acc"][:, 2], label="Left IMU Z acc")
# plt.plot(time, data["right_imu_lin_acc"][:, 2], label="Right IMU Z acc")
# plt.title("IMU Linear Acc (Z-axis)")
# plt.legend()

# # 4. 左右脚 IMU 角速度
# plt.subplot(3, 3, 4)
# plt.plot(time, data["left_imu_ang_vel"][:, 2], label="Left IMU Z ang vel")
# plt.plot(time, data["right_imu_ang_vel"][:, 2], label="Right IMU Z ang vel")
# plt.title("IMU Angular Velocity (Z-axis)")
# plt.tight_layout()
# # plt.show()

# plt.legend(),

# # 5. 脚踝高度
plt.subplot(1, 1, 1)
plt.plot(time, data["left_ankle_height"][:,2], label="Left ankle height")
plt.plot(time, data["right_ankle_height"][:,2], label="Right ankle height")
plt.plot(time, swing_standing_label[:,2], label="Swing/Standing")
plt.title("Ankle Heights")
plt.legend()



# # plt.subplot(2, 1, 2)
# # plt.plot(time, swing_standing_label[:,2], label="Swing/Standing")
# # plt.title("Ankle Heights")
# # plt.legend()

# # # 6. 关节速度
# plt.subplot(3, 3, 6)
# plt.plot(time, data["joint_vel"][:, 0], label="Left joint vel")
# plt.plot(time, data["joint_vel"][:, 2], label="Right joint vel")
# plt.title("Joint Velocities")
# plt.legend()

# plt.tight_layout()
# plt.show()
